#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/DTI.h"

#include "LeastSquare.h"

#include <teem\ell.h>

#include <omp.h>

char *fiber_pathname = NULL;
char *dst_pathname = NULL;

char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;

float fiber_length = 0.0f;

#ifndef PARAM_COUNT
#define PARAM_COUNT 17
#endif

#ifndef RESAMPLE
#define RESAMPLE 10
#endif

struct FiberParameter {
	Seed seed;
	CPoint3F parameters[PARAM_COUNT];
};

void PrintUsage()
{
	printf("Usage: ComputeNewEdgeStrength -fiber <configure_file> -dst <dst_file> -length <fiber_length>\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-length") == 0) {
			fiber_length = (float)atof(argv[++i]);
		} else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL || fiber_length < 0.0001f)
		return -2;

	return 0;
}

void FiberParameterization(CFiber *fibers, FiberParameter *parameters, const int count)
{
	int offset = int(fiber_length / stepSize * (float)RESAMPLE / (float)((PARAM_COUNT-1)/2));
	for (int i = 0; i < count; ++i) {
		parameters[i].seed = fibers[i].m_seed;
		parameters[i].parameters[PARAM_COUNT/2] = fibers[i].m_seed.pos;

		// forward
		if (fibers[i].m_fCount < 2) {
			for (int j = 0; j < (PARAM_COUNT-1)/2; ++j) {
				parameters[i].parameters[PARAM_COUNT/2+j+1] = fibers[i].m_seed.pos;
			}
		} else {
			int resamples_size = (fibers[i].m_fCount - 1) * RESAMPLE;
			CPoint3F *resamples = new CPoint3F[resamples_size];

			for (uint j = 1; j < fibers[i].m_fCount; ++j) {
				for (int k = 0; k < RESAMPLE; ++k) {
					float ratio = float(k+1) / float(RESAMPLE);
					resamples[(j-1)*RESAMPLE+k] = fibers[i].m_pF[j-1] * (1.0f - ratio) + fibers[i].m_pF[j] * ratio;
				}
			}
			for (int j = 0; j < (PARAM_COUNT-1)/2; ++j) {
				int id = (j + 1) * offset;
				id = id < resamples_size ? id : (resamples_size-1);
				parameters[i].parameters[PARAM_COUNT/2+j+1] = resamples[id];
			}
			SafeDeleteArray(resamples);
		}
		

		// backward
		if (fibers[i].m_bCount < 2) {
			for (int j = 0; j < (PARAM_COUNT-1)/2; ++j) {
				parameters[i].parameters[PARAM_COUNT/2-j-1] = fibers[i].m_seed.pos;
			}
		} else {
			int resamples_size = (fibers[i].m_bCount - 1) * RESAMPLE;
			CPoint3F *resamples = new CPoint3F[resamples_size];
			for (uint j = 1; j < fibers[i].m_bCount; ++j) {
				for (int k = 0; k < RESAMPLE; ++k) {
					float ratio = float(k+1) / float(RESAMPLE);
					resamples[(j-1)*RESAMPLE+k] = fibers[i].m_pB[j-1] * ratio + fibers[i].m_pB[j] * (1.0f - ratio);
				}
			}
			for (int j = 0; j < (PARAM_COUNT-1)/2; ++j) {
				int id = (j + 1) * offset;
				id = id < resamples_size ? id : (resamples_size-1);
				parameters[i].parameters[PARAM_COUNT/2-j-1] = resamples[id];
			}
			SafeDeleteArray(resamples);
		}
		
	}
}

//float ComputeEdgeStrength(FiberParameter *fibers, const int count)
//{
//	/* construct matrix A */
//	float *A = new float[count*3];
//	for (int i = 0; i < count; ++i) { 
//		A[i*3+0] = fibers[i].seed.pos.m_x;
//		A[i*3+1] = fibers[i].seed.pos.m_y;
//		A[i*3+2] = fibers[i].seed.pos.m_z;
//	}
//
//	float strength = 0.0f;
//
//	/* estimate the function for each parameter */
//	float *x = new float[3];
//	float *y = new float[count];
//	for (int i = 0; i < PARAM_COUNT; ++i) {
//		for (int j = 0; j < 3; ++j) {
//			for (int k = 0; k < count; ++k) { 
//				if (j == 0)
//					y[k] = fibers[k].parameters[i].m_x;
//				else if (j == 1)
//					y[k] = fibers[k].parameters[i].m_y;
//				else
//					y[k] = fibers[k].parameters[i].m_z;
//			}
//			memset(x, 0, sizeof(float)*3);
//
//			LeastSquare3(A, x, y, count, 3);
//
//			strength += x[0] * x[0] + x[1] * x[1] + x[2] * x[2];
//		}
//	}
//
//	SafeDeleteArray(A);
//	SafeDeleteArray(x);
//	SafeDeleteArray(y);
//
//	return sqrtf(strength);
//}

float ComputeEdgeStrength(FiberParameter *fibers, const int count)
{
	//float gaussian[PARAM_COUNT];
	//gaussian[PARAM_COUNT/2] = 1.0f;
	//for (int i = 0; i < PARAM_COUNT/2; ++i) {
	//	float dis = float(PARAM_COUNT / 2 - i);
	//	gaussian[i] = gaussian[PARAM_COUNT-i-1] = exp(-1.0f * dis * dis / (PARAM_COUNT * PARAM_COUNT * 0.125f));
	//}

	/* construct matrix A */
	double *A = new double[(count-1)*4];
	for (int i = 0; i < count-1; ++i) { 
		A[i*4+0] = (fibers[i+1].seed.pos.m_x - fibers[0].seed.pos.m_x) * (double)scale;
		A[i*4+1] = (fibers[i+1].seed.pos.m_y - fibers[0].seed.pos.m_y) * (double)scale;
		A[i*4+2] = (fibers[i+1].seed.pos.m_z - fibers[0].seed.pos.m_z) * (double)scale;
		A[i*4+3] = 1.0;
	}

	/* compute pseudoinverse of matrix A */
	double *pinvA = new double[4*(count-1)];
	PseudoInverse4(A, pinvA, count-1, 4);

	float strength = 0.0f;

	/* estimate the function for each parameter */
	float *x = new float[4];
	float *y = new float[count];
	for (int i = 0; i < PARAM_COUNT; ++i) {
		for (int j = 0; j < 3; ++j) {
			for (int k = 0; k < count-1; ++k) { 
				if (j == 0)
					y[k] = fibers[k+1].parameters[i].m_x;
				else if (j == 1)
					y[k] = fibers[k+1].parameters[i].m_y;
				else
					y[k] = fibers[k+1].parameters[i].m_z;
			}
			memset(x, 0, sizeof(float)*4);

			LeastSquareFast4(pinvA, x, y, count-1, 4);

			//strength += gaussian[i] * (x[0] * x[0] + x[1] * x[1] + x[2] * x[2]);
			strength += x[0] * x[0] + x[1] * x[1] + x[2] * x[2];
		}
	}

	SafeDeleteArray(A);
	SafeDeleteArray(pinvA);
	SafeDeleteArray(x);
	SafeDeleteArray(y);

	return sqrtf(strength);
}

void ReverseFiber(FiberParameter *fiber)
{
	fiber->seed.dir.m_x = -fiber->seed.dir.m_x;
	fiber->seed.dir.m_y = -fiber->seed.dir.m_y;
	fiber->seed.dir.m_z = -fiber->seed.dir.m_z;

	for (int i = 0; i < PARAM_COUNT/2; ++i) {
		//swap i with PARAM_COUNT-i-1
		CPoint3F temp = fiber->parameters[i];
		fiber->parameters[i] = fiber->parameters[PARAM_COUNT-i-1];
		fiber->parameters[PARAM_COUNT-i-1] = temp;
	}
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	/* read the configure file */
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);

	float length = fiber_length > maxLength ? maxLength : fiber_length;
	int cw, ch;

	/* compute the edge strength for each voxel */
	float *edgeStrength = new float[w*h*d];
	memset(edgeStrength, 0, sizeof(float)*w*h*d);

	CFiber *fibers = NULL;
	FiberParameter **parameters = new FiberParameter*[3];
	parameters[0] = new FiberParameter[w*h];
	parameters[1] = new FiberParameter[w*h];
	parameters[2] = new FiberParameter[w*h];

	for (int z = 0; z < 2; ++z) {
		ReadFibers(confPathname[z], vertPathname[z], &fibers, cw, ch);
		FiberParameterization(fibers, parameters[z], cw*ch);
	}
	

	int offset[9];
    offset[ 0] = 0;
    offset[ 1] = -1;
    offset[ 2] = 1;
    offset[ 3] = -w-1;
    offset[ 4] = -w;
    offset[ 5] = -w+1;
    offset[ 6] = w-1;
    offset[ 7] = w;
    offset[ 8] = w+1;

	omp_set_num_threads(64);

	for (int z = 1; z < d - 1; ++z) {
		ReadFibers(confPathname[z+1], vertPathname[z+1], &fibers, cw, ch);
		FiberParameterization(fibers, parameters[2], cw*ch);
		delete[] fibers;

		for (int y = 1; y < h - 1; ++y) {
#pragma omp parallel for            
			for (int x = 1; x < w - 1; ++x) {
				int index = y * w + x;
				FiberParameter f[27];

				for (int i = 0; i < 9; ++i) {
					memcpy(&(f[i]), &(parameters[1][index+offset[i]]), sizeof(FiberParameter));
					if (InnerProduct(parameters[1][index].seed.dir, f[i].seed.dir) < 0.0f) {
						ReverseFiber(&(f[i]));
					}
					memcpy(&(f[i+9]), &(parameters[0][index+offset[i]]), sizeof(FiberParameter));
					if (InnerProduct(parameters[1][index].seed.dir, f[i+9].seed.dir) < 0.0f) {
						ReverseFiber(&(f[i+9]));
					}
					memcpy(&(f[i+18]), &(parameters[2][index+offset[i]]), sizeof(FiberParameter));
					if (InnerProduct(parameters[1][index].seed.dir, f[i+18].seed.dir) < 0.0f) {
						ReverseFiber(&(f[i+18]));
					}
				}

				edgeStrength[z*w*h+index] = ComputeEdgeStrength(f, 27);
			}
		}

		memcpy(parameters[0], parameters[1], sizeof(FiberParameter)*w*h);
		memcpy(parameters[1], parameters[2], sizeof(FiberParameter)*w*h);
		memset(parameters[2], 0, sizeof(FiberParameter)*w*h);

		printf("z = %d\n", z);
	}

	/* save the result */
	CScalar *result = new CScalar();
	result->CreateScalar(w, h, d, edgeStrength);
	result->SaveScalarFile(dst_pathname);

	delete result;

	for (int z = 0; z < d; ++z) {
		SafeDeleteArray(confPathname[z]);
		SafeDeleteArray(vertPathname[z]);
	}
	SafeDeleteArray(confPathname);
	SafeDeleteArray(vertPathname);

	SafeDeleteArray(edgeStrength);

	SafeDeleteArray(parameters[0]);
	SafeDeleteArray(parameters[1]);
	SafeDeleteArray(parameters[2]);
	SafeDeleteArray(parameters);

	return 0;
}